from targetDriver import GoldenRunDriver, InjectionRunDriver
from datetime import datetime
import jeringa
import random
import itertools
import time
import gdb

# The RunManager abstract class, which manages general run-specific state.
# USES: ConnectionManager, TargetDriver.

class RunManager(object):
    def __init__(self, connectionManager, targetDriver):
        self.connectionManager = connectionManager
        self.targetDriver = targetDriver
        self.initialGDBSetup()

    # Clear existing breakpoints and turn off gdb confirm prompts.
    def initialGDBSetup(self):
        gdb.execute("set pagination off")
        gdb.execute("set confirm off")
        gdb.execute("set logging file /dev/null")
        gdb.execute("set logging redirect on")

    def setMuted(self, muted):
        if muted:
            gdb.execute("set logging on")
        else:
            gdb.execute("set logging off")

################################################################################

# A RunManager which manages Golden Runs.
# USES: ConnectionManager, GoldenRunDriver, EndBreakpoint.

class GoldenRunManager(RunManager):
    def __init__(self, connectionManager, SCCDetector, functionsToSkip):
        super(GoldenRunManager, self).__init__(connectionManager,
                                               GoldenRunDriver(SCCDetector, functionsToSkip))

    def startRunning(self, logFile, startPoints, endPoints):
        steps = 0
        startTime = 0
        endTime = 0

        # Connect to the target and reset it.
        self.connectionManager.connect()
        self.connectionManager.resetTarget()

        startTime = datetime.now()

        self.setMuted(True)

        print("Voy a ir al startpoint")
        self.targetDriver.advanceToStartPoint(startPoints)
        print("voy a poner los endbreakpoints")
        self.targetDriver.setEndBreakpoints(endPoints)

        with open(logFile, "w") as goldenRunLog:
            print("Estoy aca")
            steps = self.targetDriver.proceedAndLog(goldenRunLog)

        self.setMuted(False)

        endTime = datetime.now()
        deltaTime = endTime - startTime

        jeringa.log("\nFinal PC: %#x" % gdb.parse_and_eval("$pc"))
        jeringa.log("Steps: %d" % steps)
        jeringa.log("Time: %f seconds.\n" % deltaTime.total_seconds())

################################################################################

# A RunManager which manages injection runs.
# USES: ConnectionManager, InjectionRunDriver, Injector.

class InjectionRunManager(RunManager):
    def __init__(self, connectionManager, SCCDetector, functionsToSkip, injector):
        super(InjectionRunManager, self).__init__(connectionManager,
                                                  InjectionRunDriver(SCCDetector, functionsToSkip))

        self.injector = injector
        self.crashes = 0
        self.silents = 0
        self.maxSteps = 0

    def startRunning(self, logFile, numberOfRuns):
        totalTime = 0

        with open(logFile, "r") as goldenRunLog:
            # Compute the maximum number of steps we can do.
            for i, line in enumerate(goldenRunLog):
                pass

            self.maxSteps = i

            if self.maxSteps != 0:
                startTime = datetime.now()

                # Get the first logged PC.
                goldenRunLog.seek(0)
                firstPC = goldenRunLog.readline().rstrip()

                # Perform the runs.
                for i in xrange(numberOfRuns):
                    self.performInjectionRun(goldenRunLog, firstPC, i)

                endTime = datetime.now()
                totalTime = endTime - startTime

                jeringa.log("\n==========================================")
                jeringa.log("\nCrashes: %d" % self.crashes)
                jeringa.log("Tiempo total: %f" % totalTime.total_seconds())

            else:   # File is empty.
                jering.log("El archivo esta vacio.")

    def performInjectionRun(self, goldenRunLog, firstPC, iteration):
        crashDetected = False
        canPerformInjection = False

        # Go to the start of the file.
        goldenRunLog.seek(0)

        # Connect to the target, reset it, clear any remaining breakpoints,
        # and advance up to the first logged PC.
        self.connectionManager.connect()
        self.connectionManager.resetTarget()
        gdb.execute("delete")
        self.targetDriver.advanceTo("*%s" % firstPC)

        # Compute a random injection step and check if our injector
        # can actually inject at said step.
        while not canPerformInjection:
            goldenRunLog.seek(0)

            #338 #60762 #45812
            #52120
            injectionStep = random.randint(0, self.maxSteps) #347, 70, 325
            targetPC = itertools.islice(goldenRunLog, injectionStep, None).next().rstrip()
            canPerformInjection = self.injector.canInjectAt(long(targetPC, 0))

        jeringa.log("\nIteracion: %d; stepis: %d" % (iteration, injectionStep))
        jeringa.log("Target PC: %s\n" % targetPC)

        # Perform the injection.
        startTime = datetime.now()

        self.setMuted(True)

        self.targetDriver.goToInjectionStep(goldenRunLog, injectionStep, targetPC)
        injectionSuccessful = self.injector.inject()

        if injectionSuccessful:
            crashDetected = self.targetDriver.monitorAfterInjection(goldenRunLog, injectionStep)

        self.setMuted(False)

        # Disconnect from the target.
        self.connectionManager.disconnect()

        # If we detected a crash, register it.
        if crashDetected:
            self.crashes += 1

        endTime = datetime.now()
        deltaTime = endTime - startTime

        print("\nTime: %f seconds." % deltaTime.total_seconds())
        print("Crashes so far: %d" % self.crashes)
